import os
import ast
import javalang
import subprocess
import json


def identify_file_type(file_path):
    _, ext = os.path.splitext(file_path)
    return ext.lower()


def parse_python_file(file_path):
    with open(file_path, "r", encoding="utf-8") as file:
        source_code = file.read()
    tree = ast.parse(source_code)
    return tree


def parse_java_file(file_path):
    with open(file_path, "r", encoding="utf-8") as file:
        source_code = file.read()
    tree = javalang.parse.parse(source_code)
    return tree


def parse_go_file(file_path):
    # 使用 go/parser 和 go/token 来解析Go文件
    cmd = f"parser.exe {file_path}"
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return result.stdout


def ast_to_json(tree):
    """将AST转换为JSON格式"""

    def node_to_dict(node):
        if isinstance(node, ast.AST):
            return {
                "type": type(node).__name__,
                "fields": {k: node_to_dict(v) for k, v in ast.iter_fields(node)}
            }
        elif isinstance(node, list):
            return [node_to_dict(item) for item in node]
        else:
            return str(node)

    return json.dumps(node_to_dict(tree), indent=4)


def java_ast_to_dict(node):
    """将Java AST节点转换为字典"""
    if isinstance(node, javalang.tree.CompilationUnit):
        return {
            "type": "CompilationUnit",
            "classes": [class_to_dict(c) for c in node.types]
        }
    else:
        return {"type": type(node).__name__, "details": str(node)}


def class_to_dict(node):
    """将Java类节点转换为字典"""
    return {
        "type": "ClassDeclaration",
        "name": node.name,
        "methods": [method_to_dict(m) for m in node.body if isinstance(m, javalang.tree.MethodDeclaration)]
    }


def method_to_dict(node):
    """将Java方法节点转换为字典"""
    return {
        "type": "MethodDeclaration",
        "name": node.name,
        "statements": [statement_to_dict(s) for s in node.body]
    }


def statement_to_dict(node):
    """将Java语句节点转换为字典"""
    if isinstance(node, javalang.tree.StatementExpression):
        return {
            "type": "StatementExpression",
            "expression": expression_to_dict(node.expression)
        }
    return {"type": type(node).__name__, "details": str(node)}


def expression_to_dict(node):
    """将Java表达式节点转换为字典"""
    if isinstance(node, javalang.tree.MethodInvocation):
        return {
            "type": "MethodInvocation",
            "member": node.member,
            "arguments": [expression_to_dict(a) for a in node.arguments]
        }
    elif isinstance(node, javalang.tree.Literal):
        return {
            "type": "Literal",
            "value": node.value
        }
    return {"type": type(node).__name__, "details": str(node)}


def save_json_to_file(json_output, output_file_path):
    """将JSON输出保存到指定的文本文件中"""
    with open(output_file_path, 'w', encoding='utf-8') as file:
        file.write(json_output)


def print_json_output(tree, language, output_file_path):
    if language == ".py":
        json_output = ast_to_json(tree)
    elif language == ".java":
        json_output = json.dumps(java_ast_to_dict(tree), indent=4)
    elif language == ".go":
        json_output = tree  # 假设Go AST已经是JSON格式
    else:
        print("不支持的文件类型")
        return

    save_json_to_file(json_output, output_file_path)


def main(file_path, output_file_path):
    file_type = identify_file_type(file_path)

    if file_type == ".py":
        tree = parse_python_file(file_path)
    elif file_type == ".java":
        tree = parse_java_file(file_path)
    elif file_type == ".go":
        tree = parse_go_file(file_path)
    else:
        print("不支持的文件类型")
        return

    print_json_output(tree, file_type, output_file_path)


def final_AST_JSON_get(path):
    output_file_path = 'data.json'  # AST输出的文件路径
    main(path, output_file_path)

if __name__ == "__main__":
    path='muma.py'
    final_AST_JSON_get(path)